// (C) 2022 Intel Corporation. All rights reserved.
// Your use of Intel Corporation's design tools, logic functions and other 
// software and tools, and its AMPP partner logic functions, and any output 
// files from any of the foregoing (including device programming or simulation 
// files), and any associated documentation or information are expressly subject 
// to the terms and conditions of the Intel Program License Subscription 
// Agreement, Intel FPGA IP License Agreement, or other applicable 
// license agreement, including, without limitation, that your use is for the 
// sole purpose of programming logic devices manufactured by Intel and sold by 
// Intel or its authorized distributors.  Please refer to the applicable 
// agreement for further details.

module A_cse_ocs_dpa_aes10c_key (
  
  input  logic         sys_clk,              // system clock
  input  logic         reset_n,              // system reset
  input  logic [255:0] key_in,               // Key-Input Bus from system
  input  logic         encrypt,              // cypher operation - 1 => encrypt, 0 => decrypt
  input  logic         sbox_input_sel,       // 1 => select G xor H as SBOX input, 0 => select H
  input  logic         store_key,            // store the key
  input  logic [7:0]   rcon,                 // lower 4-bits of the round constant
  input  logic         aes_key_start,            // Start the AES operation
  input  logic         flush,                // Flush flops
  input  logic         even_rounds,
  input  logic         odd_rounds,
  input  logic [1:0]   i_key_size,
  output logic [127:0] round_key,            // four words of the the working key provided to the state module
  output logic [255:0] aes_expand_key    
);

`include "A_cse_ocs_dpa_aes10c_localparams.v"
/**********************************************************************/
// Data Type Declarations (Signal Declarations)
/**********************************************************************/
   wire [31:0] keyA,keyB,keyC,keyD,keyE,keyF,keyG,keyH;     // key words that comprise the working key
   wire [31:0] A_xor_B;                 // A+B
   wire [31:0] C_xor_D;                 // C+D
   wire [31:0] E_xor_F;                 // E+F
   wire [31:0] G_xor_H;                 // G+H
   wire [31:0] sbox_xor_rcon;           // SBOX output + round constant
   wire [31:0] sbox_input;              // input to the SBOX input mux
   wire [31:0] sbox_input1;              // input to the SBOX input mux
   wire [31:0] sbox_output;             // sbox output
   wire [31:0] sbox_output1;             // sbox output
   wire [31:0] rotword_input;           // input to the rotword operation

   reg  [255:0] working_key;             // 256-bit working key register

   logic   [31:0] next_A;               // computed keyword A - all keys
   logic   [31:0] next_B;               // computed keyword B - all keys
   logic   [31:0] next_C;               // computed keyword C - all keys
   logic   [31:0] next_D;               // computed keyword D - all keys
   logic   [31:0] next_E;               // computed keyword E - all keys
   logic   [31:0] next_F;               // computed keyword F - all keys
   logic   [31:0] next_G;               // computed keyword G - all keys
   logic   [31:0] next_H;               // computed keyword H - all keys

//**********************************************************************/
// Main Code
//**********************************************************************/

  assign aes_expand_key = ~i_key_size[1]? {128'b0,next_E,next_F,next_G,next_H} : 
					  {keyE,keyF,keyG,keyH,next_A,next_B,next_C,next_D};

  // 32b Key-word assignments
  assign keyA = working_key[255:224];
  assign keyB = working_key[223:192];
  assign keyC = working_key[191:160];
  assign keyD = working_key[159:128];
  assign keyE = working_key[127:96 ];
  assign keyF = working_key[95:64  ];
  assign keyG = working_key[63:32  ];
  assign keyH = working_key[31:0   ];

  assign round_key = (i_key_size == NK_128) ? {keyE,keyF,keyG,keyH} :
		    ((i_key_size == NK_256) && encrypt && odd_rounds) | ((i_key_size == NK_256) && ~encrypt && even_rounds) ? {keyE,keyF,keyG,keyH} :
		    {keyA,keyB,keyC,keyD};

  // Intermediate terms used in key expansion (32b wide)
  assign A_xor_B = keyA ^ keyB;
  assign C_xor_D = keyC ^ keyD;
  assign E_xor_F = keyE ^ keyF;
  assign G_xor_H = keyG ^ keyH;

  // Select the input keyword for rotate and sBox functions
  assign rotword_input = (i_key_size == NK_256) ? (sbox_input_sel ? keyD : keyH) :
						   (sbox_input_sel ? G_xor_H : keyH);

  // Assign sBox input
  assign sbox_input = rotword(rotword_input);

  assign sbox_xor_rcon = sbox_output ^ {rcon,24'h000000};
  assign sbox_input1 = ((i_key_size == NK_256) && ~encrypt) ? G_xor_H :
							       sbox_xor_rcon ^ A_xor_B ^ C_xor_D ;

  // S-Box substitutions. As defined in the AES standard, this is the GF256 inverse
  // followed by an affine transformation.  These are byte-wide operations so four
  // are required to process the 32-bit half keywords.

  // Use optimized SBOX module which maps the inputs from GF(2^8) to GF(2^4)^2
  // and computes the inverse in GF(2^4). The final SBOX output is mapped back
  // to GF(2^8) for subsequent computations.

  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox_3 ( .in(sbox_input[31:24]), .out(sbox_output[31:24]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox_2 ( .in(sbox_input[23:16]), .out(sbox_output[23:16]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox_1 ( .in(sbox_input[15:8 ]), .out(sbox_output[15:8 ]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox_0 ( .in(sbox_input[7:0  ]), .out(sbox_output[7:0  ]));

  // are required to process the 32-bit half keywords for AES-256
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox1_3 ( .in(sbox_input1[31:24]), .out(sbox_output1[31:24]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox1_2 ( .in(sbox_input1[23:16]), .out(sbox_output1[23:16]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox1_1 ( .in(sbox_input1[15:8 ]), .out(sbox_output1[15:8 ]));
  A_cse_ocs_dpa_aes10c_gf16_sq_sbox_withmap sbox1_0 ( .in(sbox_input1[7:0  ]), .out(sbox_output1[7:0  ]));

  always_comb
    begin
	if ((i_key_size == NK_256) & encrypt)         //    Encryption
	   begin
            next_A = sbox_xor_rcon ^ keyA;             // SBOX(H)+RCON+A       = next A
            next_B = sbox_xor_rcon ^ A_xor_B;          // SBOX(H)+RCON+A+B     = next B
            next_C = next_B    ^ keyC;                 // SBOX(H)+RCON+A+B+C   = next C 
	    next_D = next_B    ^ C_xor_D;              // SBOX(H)+RCON+A+B+C+D = next D
            next_E = sbox_output1 ^ keyE;               // SBOX(D)+E            = next E
            next_F = sbox_output1 ^ E_xor_F;            // SBOX(D)+E+F          = next F
            next_G = next_F      ^ keyG;               // SBOX(D)+E+F+G        = next G
            next_H = next_F      ^ G_xor_H;            // SBOX(D)+E+F+G+H      = next H
	   end                                                   

	else if((i_key_size == NK_256) & ~encrypt)    //    Decryption
           begin
            next_A = sbox_output1 ^ keyA;             	// [SBOX(H)+RCON+A]       - [SBOX(H)+RCON]       = previous A
            next_B = A_xor_B;                          // [SBOX(H)+RCON+A+B]     - [SBOX(H)+RCON+A]     = previous B
            next_C = keyB ^ keyC;                      // [SBOX(H)+RCON+A+B+C]   - [SBOX(H)+RCON+A+B]   = previous C
            next_D = C_xor_D;                          // [SBOX(H)+RCON+A+B+C+D] - [SBOX(H)+RCON+A+B+C] = previous D
            next_E = sbox_xor_rcon ^ keyE;             // [SBOX(D)+E]            - [SBOX(D)]            = previous E
            next_F = E_xor_F;                          // [SBOX(D)+E+F]          - [SBOX(D)+E]          = previous F
            next_G = keyF ^ keyG;                      // [SBOX(D)+E+F+G]        - [SBOX(D)+E+F]        = previous G
            next_H = G_xor_H;                          // [SBOX(D)+E+F+G+H]      - [SBOX(D)+E+F+G]      = previous H
           end

        else if ((i_key_size == NK_128) & encrypt)          //  Encryption 128
           begin
            next_A = 32'b0;                          // [SBOX(H)+RCON+A]       - [SBOX(H)+RCON]       = previous A
            next_B = 32'b0;                          // [SBOX(H)+RCON+A+B]     - [SBOX(H)+RCON+A]     = previous B
            next_C = 32'b0;                          // [SBOX(H)+RCON+A+B+C]   - [SBOX(H)+RCON+A+B]   = previous C
            next_D = 32'b0;
            next_E = sbox_xor_rcon ^ keyE;           // SBOX(H)+RCON+E       = next E
            next_F = sbox_xor_rcon ^ E_xor_F;        // SBOX(H)+RCON+E+F     = next F
            next_G = next_F    ^ keyG;               // SBOX(H)+RCON+E+F+G   = next G
            next_H = next_F    ^ G_xor_H;            // SBOX(H)+RCON+E+F+G+H = next H
           end

        else if((i_key_size == NK_128) & ~encrypt)    //  Decryption 128
           begin
            next_A = 32'b0;               	// [SBOX(H)+RCON+A]       - [SBOX(H)+RCON]       = previous A
            next_B = 32'b0;                     // [SBOX(H)+RCON+A+B]     - [SBOX(H)+RCON+A]     = previous B
            next_C = 32'b0;                      // [SBOX(H)+RCON+A+B+C]   - [SBOX(H)+RCON+A+B]   = previous C
            next_D = 32'b0;
            next_E = sbox_xor_rcon ^ keyE;       // [SBOX(H)+RCON+E]       - [SBOX(H-G)+RCON]     = previous E
            next_F = E_xor_F;                    // [SBOX(H)+RCON+E+F]     - [SBOX(H)+RCON+E]     = previous F
            next_G = keyF  ^ keyG;               // [SBOX(H)+RCON+E+F+G]   - [SBOX(H)+RCON+E+F]   = previous G
            next_H = G_xor_H;                    // [SBOX(H)+RCON+E+F+G+H] - [SBOX(H)+RCON+E+F+G] = previous H
           end // if ((i_key_size == NK_128) & ~encrypt)
              
	else
           begin
            next_A = 32'b0;               	// [SBOX(H)+RCON+A]       - [SBOX(H)+RCON]       = previous A
            next_B = 32'b0;                     // [SBOX(H)+RCON+A+B]     - [SBOX(H)+RCON+A]     = previous B
            next_C = 32'b0;                      // [SBOX(H)+RCON+A+B+C]   - [SBOX(H)+RCON+A+B]   = previous C
            next_D = 32'b0;
            next_E = sbox_xor_rcon ^ keyE;       // [SBOX(H)+RCON+E]       - [SBOX(H-G)+RCON]     = previous E
            next_F = E_xor_F;                    // [SBOX(H)+RCON+E+F]     - [SBOX(H)+RCON+E]     = previous F
            next_G = keyF  ^ keyG;               // [SBOX(H)+RCON+E+F+G]   - [SBOX(H)+RCON+E+F]   = previous G
            next_H = G_xor_H;                    // [SBOX(H)+RCON+E+F+G+H] - [SBOX(H)+RCON+E+F+G] = previous H
	   end
    end


  // Internal registers holding the working key
  always_ff @ (posedge sys_clk or negedge reset_n)
    begin: internal_registers
       if (~reset_n)
         begin
           working_key     <= {256{1'b0}};
         end
       else
         begin

           if (flush) begin
              working_key <= 256'h0;
           end else begin
              if (aes_key_start) begin
                 working_key <= key_in;
              end else if (store_key) begin
                working_key[255:224]   <= next_A;
                working_key[223:192]   <= next_B;
                working_key[191:160]   <= next_C;
                working_key[159:128]   <= next_D;
                working_key[127:96]    <= next_E;
                working_key[95:64]     <= next_F;
                working_key[63:32]     <= next_G;
                working_key[31:0]      <= next_H;
              end
           end
         end
    end // block: internal_registers

    /**********************************************************************/
    //  Functions used for key processing
    /**********************************************************************/

    /****************************************************************/
    // Function name: rotword
    // Inputs       : input_word[31:0]
    // Description  : Performs a cyclic permutation (left shift)
    //                on the bytes in the input_word:
    //                {b0,b1,b2,b3} => {b1,b2,b3,b0}.
    //                Result is a word.
    /****************************************************************/
    function automatic [31:0] rotword;
      input [31:0] input_word;
      begin
        rotword = {input_word[23:16],
                   input_word[15:8],
                   input_word[7:0],
                   input_word[31:24]};
      end
    endfunction

endmodule
